"""
该脚本用于将本地数据保存到向量库中
本地数据格式为jsonl
{
    "input":"问题",
    "output":"答案",
    "id":"唯一标识"
}
"""
import sys
sys.path.append('/opt/data/private/liuteng/code/OpenBA-3B')

import qdrant_client
import tqdm
from data_loader.json_data_loader import ChatJSONLLoader
from langchain.vectorstores.qdrant import Qdrant
from langchain.schema import Document
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from qdrant_client.http.models import Distance, VectorParams

QDRANT_URL='10.10.80.174:6333'
EMBEDING_MODLE="/opt/data/private/liuteng/model/bge-large-zh"

chatJSONLLoader=ChatJSONLLoader(
    file_path="data/merge_data.jsonl",
    encoding="utf-8",
    id_key="id",
    question_key="input",
    answer_key="output"
)

print("开始加载embeding模型")

embeddings=HuggingFaceEmbeddings(
    model_name=EMBEDING_MODLE
)




if __name__=="__main__":
    documents=chatJSONLLoader.load()
    
    client = qdrant_client.QdrantClient(
        url=QDRANT_URL,
        prefer_grpc=False,
    )

    client.create_collection(
        collection_name="validate_data",
        vectors_config=VectorParams(size=1024, distance=Distance.COSINE)
    )

    qdrant = Qdrant(
        client=client,
        collection_name="validate_data",
        embeddings=embeddings,
    )

    for document in tqdm.tqdm(documents):
        qdrant.add_documents(documents=[document])


    
    print("向量化成功")



